import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.family'] = 'SimHei'  # Windows系统常用中文字体

# 参数设置
Pop_size = 10
Max_T = 75
Dim = 2
Num_obstacles = 50
safe_distance = 5
bounds = np.array([[0, 100], [0, 100]])


# 随机生成起点和终点
def generate_start_and_goal(bounds, safe_distance,
                            min_sep=None, margin=5):
    """
    保证起点在左上矩形，终点在右下矩形，且两者欧氏距离 >= min_sep
    bounds:  np.array([[xmin, xmax], [ymin, ymax]])
    margin:  距离边界的最小留白，防止起点/终点贴边
    min_sep: 两点间最小距离，默认等于 safe_distance
    """
    if min_sep is None:
        min_sep = safe_distance

    x_min, x_max = bounds[0, 0], bounds[0, 1]
    y_min, y_max = bounds[1, 0], bounds[1, 1]

    # 计算中心点，把地图切成左上 / 右下两个矩形
    x_c = (x_min + x_max) / 2
    y_c = (y_min + y_max) / 2

    # 左上矩形边界
    left_x1, right_x1 = x_min + margin, x_c - margin
    bottom_y1, top_y1    = y_c + margin, y_max - margin

    # 右下矩形边界
    left_x2, right_x2 = x_c + margin, x_max - margin
    bottom_y2, top_y2    = y_min + margin, y_c - margin

    # 确保矩形有效
    assert left_x1 < right_x1 and bottom_y1 < top_y1, "左上区域过小，请调大地图或减小 margin"
    assert left_x2 < right_x2 and bottom_y2 < top_y2, "右下区域过小，请调大地图或减小 margin"

    while True:
        # 在左上矩形采起点
        start = np.array([np.random.uniform(left_x1, right_x1),
                          np.random.uniform(bottom_y1, top_y1)])
        # 在右下矩形采终点
        goal  = np.array([np.random.uniform(left_x2, right_x2),
                          np.random.uniform(bottom_y2, top_y2)])

        if np.linalg.norm(start - goal) >= min_sep:
            return start, goal


# 随机生成障碍物
def generate_obstacles(num_obstacles, bounds, start, goal, safe_distance):
	obstacles = []
	while len(obstacles) < num_obstacles:
		obs = np.random.uniform(bounds[:, 0], bounds[:, 1], size=(1, 2)).flatten()
		if (np.linalg.norm(obs - start) >= safe_distance and
				np.linalg.norm(obs - goal) >= safe_distance and
				all(np.linalg.norm(obs - other) >= 2 * safe_distance for other in obstacles)):
			obstacles.append(obs)
	return np.array(obstacles)


# 适应度函数（标准GWO与改进GWO共用）
def fitness_function(position, goal, obstacles, safe_distance):
	distance_to_goal = np.linalg.norm(position - goal)
	penalty = sum([1000 / dis_to_obs for obs in obstacles
				   if (dis_to_obs := np.linalg.norm(position - obs)) < safe_distance])
	return distance_to_goal + penalty


# 标准GWO算法实现
def standard_gwo(Pop_size, Max_T, Dim, start, goal, bounds, obstacles, safe_distance):
	positions = np.random.uniform(bounds[:, 0], bounds[:, 1], size=(Pop_size, Dim))
	alpha_pos, beta_pos, delta_pos = start.copy(), None, None
	alpha_score, beta_score, delta_score = float('inf'), float('inf'), float('inf')
	best_path, best_scores = [start.copy()], []

	for iter in range(Max_T):
		for i in range(Pop_size):
			fitness = fitness_function(positions[i, :], goal, obstacles, safe_distance)

			if fitness < alpha_score:
				delta_score, delta_pos = beta_score, beta_pos
				beta_score, beta_pos = alpha_score, alpha_pos
				alpha_score, alpha_pos = fitness, positions[i, :].copy()
			elif fitness < beta_score:
				delta_score, delta_pos = beta_score, beta_pos
				beta_score, beta_pos = fitness, positions[i, :].copy()
			elif fitness < delta_score:
				delta_score, delta_pos = fitness, positions[i, :].copy()

		best_path.append(alpha_pos.copy())
		best_scores.append(alpha_score)

		for i in range(Pop_size):
			for j in range(Dim):
				r1_alpha, r2_alpha = np.random.rand(), np.random.rand()
				A1 = 2 * (2 - iter * (2 / Max_T)) * r1_alpha - (2 - iter * (2 / Max_T))
				C1 = 2 * r2_alpha
				D_alpha = abs(C1 * alpha_pos[j] - positions[i, j])
				X1 = alpha_pos[j] - A1 * D_alpha

				A2, C2 = 2 * (2 - iter * (2 / Max_T)) * np.random.rand() - (
						2 - iter * (2 / Max_T)), 2 * np.random.rand()
				D_beta = abs(C2 * beta_pos[j] - positions[i, j]) if beta_pos is not None else 0
				X2 = beta_pos[j] - A2 * D_beta

				A3, C3 = 2 * (2 - iter * (2 / Max_T)) * np.random.rand() - (
						2 - iter * (2 / Max_T)), 2 * np.random.rand()
				D_delta = abs(C3 * delta_pos[j] - positions[i, j]) if delta_pos is not None else 0
				X3 = delta_pos[j] - A3 * D_delta

				positions[i, j] = (X1 + X2 + X3) / 3

		positions = np.clip(positions, bounds[:, 0], bounds[:, 1])

	return best_path, best_scores, alpha_pos


# 改进GWO算法实现
def modified_gwo(Pop_size, Max_T, Dim, start, goal, bounds, obstacles, safe_distance):
	positions = np.random.uniform(bounds[:, 0], bounds[:, 1], size=(Pop_size, Dim))
	alpha_pos, beta_pos, delta_pos = start.copy(), None, None
	alpha_score, beta_score, delta_score = float('inf'), float('inf'), float('inf')
	best_path, best_scores = [start.copy()], []

	for iter in range(Max_T):
		a = 2 - iter * (2 / Max_T)

		for i in range(Pop_size):
			fitness = fitness_function(positions[i, :], goal, obstacles, safe_distance)
			if fitness < alpha_score:
				delta_score, delta_pos = beta_score, beta_pos
				beta_score, beta_pos = alpha_score, alpha_pos
				alpha_score, alpha_pos = fitness, positions[i, :].copy()
			elif fitness < beta_score:
				delta_score, delta_pos = beta_score, beta_pos
				beta_score, beta_pos = fitness, positions[i, :].copy()
			elif fitness < delta_score:
				delta_score, delta_pos = fitness, positions[i, :].copy()

		best_path.append(alpha_pos.copy())
		best_scores.append(alpha_score)

		for i in range(Pop_size):
			for j in range(Dim):
				r1, r2 = np.random.rand(), np.random.rand()
				A1, C1 = 2 * a * r1 - a, 2 * r2
				D_alpha = abs(C1 * alpha_pos[j] - positions[i, j])
				X1 = alpha_pos[j] - A1 * D_alpha

				A2, C2 = 2 * a * np.random.rand() - a, 2 * np.random.rand()
				D_beta = abs(C2 * beta_pos[j] - positions[i, j]) if beta_pos is not None else 0
				X2 = beta_pos[j] - A2 * D_beta

				A3, C3 = 2 * a * np.random.rand() - a, 2 * np.random.rand()
				D_delta = abs(C3 * delta_pos[j] - positions[i, j]) if delta_pos is not None else 0
				X3 = delta_pos[j] - A3 * D_delta

				positions[i, j] = (0.6 * X1 + 0.3 * X2 + 0.1 * X3)

		positions = np.clip(positions, bounds[:, 0], bounds[:, 1])

	return best_path, best_scores, alpha_pos


start, goal = generate_start_and_goal(bounds, safe_distance)
obstacles = generate_obstacles(Num_obstacles, bounds, start, goal, safe_distance)

# 运行两种GWO算法
standard_path, standard_scores, standard_final_pos = standard_gwo(Pop_size, Max_T, Dim, start, goal, bounds, obstacles,
																  safe_distance)
modified_path, modified_scores, modified_final_pos = modified_gwo(Pop_size, Max_T, Dim, start, goal, bounds, obstacles,
																  safe_distance)

# 绘制最终路径
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.title("标准GWO - 最终路径")
plt.plot(start[0], start[1], 'go', markersize=10, label='起点')
plt.plot(goal[0], goal[1], 'ro', markersize=10, label='终点')
plt.plot(obstacles[:, 0], obstacles[:, 1], 'ks', markersize=4, label='障碍物')
standard_path_np = np.array(standard_path)
plt.plot(standard_path_np[:, 0], standard_path_np[:, 1], 'c-', linewidth=1.5, label='路径')
plt.scatter(standard_path_np[:, 0], standard_path_np[:, 1], color='c', s=20, marker='o')
plt.plot(standard_final_pos[0], standard_final_pos[1], 'c*', markersize=12, label='最终Alpha位置')
plt.legend()

plt.subplot(1, 2, 2)
plt.title("改进GWO - 最终路径")
plt.plot(start[0], start[1], 'go', markersize=10, label='起点')
plt.plot(goal[0], goal[1], 'ro', markersize=10, label='终点')
plt.plot(obstacles[:, 0], obstacles[:, 1], 'ks', markersize=4, label='障碍物')
modified_path_np = np.array(modified_path)
plt.plot(modified_path_np[:, 0], modified_path_np[:, 1], 'm-', linewidth=1.5, label='路径')
plt.scatter(modified_path_np[:, 0], modified_path_np[:, 1], color='m', s=20, marker='o')
plt.plot(modified_final_pos[0], modified_final_pos[1], 'm*', markersize=12, label='最终Alpha位置')
plt.legend()
plt.show()

# 绘制收敛曲线对比
plt.figure()
plt.plot(range(Max_T), standard_scores, 'b-', label='标准GWO')
plt.plot(range(Max_T), modified_scores, 'r-', label='改进GWO')
plt.xlabel('迭代次数')
plt.ylabel('最优适应度值')
plt.title('收敛曲线对比')
plt.legend()
plt.show()

# 计算最优路径长度
standard_path_Len = sum(np.linalg.norm(standard_path_np[i] - standard_path_np[i + 1])
						for i in range(len(standard_path_np) - 1))

print(f'标准路径长度: {standard_path_Len}')

modified_path_Len = sum(np.linalg.norm(modified_path_np[i] - modified_path_np[i + 1])
						for i in range(len(modified_path_np) - 1))

print(f'改进路径长度: {modified_path_Len}')
